from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型

from tool import get_azure_endpoint, get_api_key, get_api_version

if __name__ == '__main__':
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    sentence1_chinese = "我喜欢狗"
    sentence2_chinese = "我喜欢犬科动物"
    sentence3_chinese = "外面的天气很糟糕"
    embedding1_chinese = embedding.embed_query(sentence1_chinese)
    embedding2_chinese = embedding.embed_query(sentence2_chinese)
    embedding3_chinese = embedding.embed_query(sentence3_chinese)

    import numpy as np

    score1 = np.dot(embedding1_chinese, embedding2_chinese)
    print(score1)

    score2 = np.dot(embedding1_chinese, embedding3_chinese)
    print(score2)

    score3 = np.dot(embedding2_chinese, embedding3_chinese)
    print(score3)